package com.fy.fyai.common.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 多模型管理器
 * 统一管理不同的 ChatModel，支持动态切换
 *
 * @author fwq
 * @date 2025-10-26
 */
@Slf4j
@Component
public class MultiModelManager {

    private final Map<String, ChatModel> modelMap = new HashMap<>();

    /**
     * 默认使用的模型名称
     */
    public static final String DEFAULT_MODEL = "ollama";
    public static final String OLLAMA_MODEL = "ollama";
    public static final String OPENAI_MODEL = "openai";

    @Autowired
    public MultiModelManager(@Qualifier("ollamaChatModel") ChatModel ollamaChatModel,
                            @Qualifier("openAiChatModel") ChatModel openAiChatModel) {
        // 注册所有可用的模型
        modelMap.put(OLLAMA_MODEL, ollamaChatModel);
        modelMap.put(OPENAI_MODEL, openAiChatModel);

        log.info("多模型管理器初始化完成");
        log.info("可用模型: {}", modelMap.keySet());
        log.info("默认模型: {}", DEFAULT_MODEL);
    }

    /**
     * 获取指定模型
     *
     * @param modelName 模型名称 (ollama 或 openai)
     * @return ChatModel
     */
    public ChatModel getModel(String modelName) {
        if (modelName == null || modelName.isEmpty()) {
            modelName = DEFAULT_MODEL;
        }

        ChatModel model = modelMap.get(modelName.toLowerCase());
        if (model == null) {
            log.warn("模型{}不存在，使用默认模型{}", modelName, DEFAULT_MODEL);
            model = modelMap.get(DEFAULT_MODEL);
        }
        return model;
    }

    /**
     * 获取默认模型
     *
     * @return ChatModel
     */
    public ChatModel getDefaultModel() {
        return getModel(DEFAULT_MODEL);
    }

    /**
     * 获取 Ollama 模型
     *
     * @return ChatModel
     */
    public ChatModel getOllamaModel() {
        return getModel(OLLAMA_MODEL);
    }

    /**
     * 获取 OpenAI 模型
     *
     * @return ChatModel
     */
    public ChatModel getOpenAiModel() {
        return getModel(OPENAI_MODEL);
    }

    /**
     * 检查模型是否存在
     *
     * @param modelName 模型名称
     * @return true if exists
     */
    public boolean hasModel(String modelName) {
        return modelMap.containsKey(modelName.toLowerCase());
    }

    /**
     * 获取所有可用的模型名称
     *
     * @return 模型名称集合
     */
    public Set<String> getAvailableModels() {
        return modelMap.keySet();
    }
}
